#!/usr/bin/env python
# -*- coding=utf-8 -*-
import re
from time import time

import requests
import bs4
import colorama

from codon_aa import codon_tab

def make_colors():
    """
    make some colors in command line
    """
    colors = {}
    colorama.init()
    colors['red'] = colorama.Fore.RED
    colors['green'] = colorama.Fore.GREEN
    colors['green_ex'] = colorama.Fore.LIGHTGREEN_EX
    colors['white'] = colorama.Fore.LIGHTWHITE_EX
    colors['magenta_ex'] = colorama.Fore.LIGHTMAGENTA_EX
    colors['magenta'] = colorama.Fore.MAGENTA
    colors['cyan'] = colorama.Fore.CYAN
    colors['cyan_ex'] = colorama.Fore.LIGHTCYAN_EX
    colors['yellow'] = colorama.Fore.YELLOW
    colors['bright'] = colorama.Style.BRIGHT
    colors['back'] = colorama.Style.RESET_ALL
    return colors

#Global Var
basic_url = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi?'
colors = make_colors()

def get_NM_info(NM_id,try_time=3):
    """
    get mRNA info from NCBI and extract CDS seq and amino acid seq
    """
    search_url = basic_url + 'db=nuccore&id=' + NM_id + '&retmode=xml'
    try:
        resp = requests.get(search_url)
    except Exception as e:
        if try_time < 0:
            print(e)
            exit(0)
        else:
            print('try again {try_time}'.format(try_time=3-try_time))
            return get_NM_info(NM_id,try_time-1)
    if resp:
        return resp.text
    else:
        raise Exception('Get Nothing from NCBI!')    

def parse_xml(xml_tree):
    """
    get CDS region and its translation 
    """
    soup = bs4.BeautifulSoup(xml_tree,'xml')
    seq = soup.select('GBSeq_sequence')[0].text
    cds_loc,translation = find_range(xml_tree)
    # print(cds_loc)
    # print(translation)
    cds_seq = seq[cds_loc[0]-1:cds_loc[1]] #cds_loc是1-base的，因此这里减去1
    utr3 = seq[cds_loc[1]:]
    # print(cds_seq)
    assert len(cds_seq)%3 == 0,"cds_seq is not fold of 3!"
    return cds_seq,translation,utr3

def find_range(xml_tree):
    """
    I just can not find CDS location by bs4,so
    this function is the replaced method
    """
    xml_tree_list = xml_tree.split('\n')
    for idx,line in enumerate(xml_tree_list):
        if "<GBFeature_key>CDS</GBFeature_key>" in line:
            location = re.findall(r'\d+',xml_tree_list[idx+1])
            location = [int(l) for l in location]
            # print(location)
        
        if "<GBQualifier_name>translation" in line:
            translation = re.search(r'<GBQualifier_value>(.*?)</GBQualifier_value>',xml_tree_list[idx+1])
            if translation is not None:
                translation = translation.group(1)
                # print(translation)
            else:
                raise Exception('Can NOT find translation!')

    if (not locals().get("translation")) or (not locals().get("location")):
        raise Exception('Can not find translation or cds location in xml tree!')

    return location,translation

class ShowVar():

    def __init__(self,cdna,cds_seq,translation,utr3,num_of_each_line=10):
        self.cdna = cdna
        self.utr3 = utr3
        self.cds_seq = cds_seq
        self.translation = translation
        self.num_of_each_line = num_of_each_line
    
    def parse_cdna(self):
        """
        split the cdna and parse its mutation site
        built the new sequence
        """
        re_del = re.search(r'^c.(\d+)_(\d+)del$',self.cdna)
        re_single_del = re.search(r'^c.(\d+)del([ATCG])$',self.cdna)
        re_dup = re.search(r'^c.(\d+)dup([ATCG])$',self.cdna)
        re_single_replace = re.search(r'^([ATCG])(\d+)([ATCG])$',self.cdna)
        seq_list = [b for b in self.cds_seq]
        if re_single_replace is not None:
            raw_base,position,replace_base = re_single_replace.groups()
            position = int(position)
            pos_idx = position-1
            ref_base = seq_list[pos_idx]
            if ref_base == raw_base.lower():
                seq_list[pos_idx] = replace_base
            else:
                raise Exception('raw base is unmatched with reference base at {position},raw base:{raw_base},ref base:{ref_base}'.format(
            **locals()))
        elif re_dup is not None:
            position,insert_base = re_dup.groups()
            position = int(position)
            pos_idx = position-1
            ref_base = seq_list[pos_idx]
            seq_list.insert(pos_idx,insert_base)

        elif re_del is not None:
            position,end_position = re_del.groups() #闭区间
            position = int(position)
            end_position = int(end_position)
            seq_list = seq_list[:position-1] + seq_list[end_position:]
            #这里end_position用于index需要先减1，然后因为是闭区间所以需要加1，前后抵消因此是end_posiiton-1+1
        elif re_single_del is not None:
            position,raw_base = re_single_del.groups()
            position = int(position)
            pos_idx = position-1
            ref_base = seq_list[pos_idx]
            if ref_base == raw_base.lower():
                del seq_list[pos_idx]
            else:
                raise Exception('raw base is unmatched with reference base at {position},raw base:{raw_base},ref base:{ref_base}'.format(
            **locals()))
        else:
            raise Exception('unsupport cDNA format {cdna}'.format(cdna=self.cdna))
        
        #justify position by fold of 3
        if position%3 == 1:
            justify_position = position
        elif position%3 == 2:
            justify_position = position -1
        elif position%3 == 0:
            justify_position = position -2

        return "".join(seq_list),justify_position
    
    def extract_cds_seq(self,pos_list):
        """
        """
        cds_codon_list = []
        cds_aa_list = []
        for idx in range(len(pos_list)):
            start = pos_list[idx] 
            start_idx = start -1 # pos_list中的数字都是1base的，因此这里要减一
            aa_idx = int((start + 2 )/3 -1)
            try:
                end = pos_list[idx+1]
                end_idx = end - 1 
            except IndexError:
                cds_codon_list.append(self.cds_seq[start_idx:start_idx+3])
                cds_aa_list.append(self.translation[aa_idx])
                return cds_codon_list,cds_aa_list

            cds_codon_list.append(self.cds_seq[start_idx:end_idx])
            cds_aa_list.append(self.translation[aa_idx])
        return cds_codon_list,cds_aa_list

    def show_cdna_main(self):
        """
        show cdna
        """
        seq,pos = self.parse_cdna()
        #pos是1base的，因此用作index需要减1
        cdna_list,aa_list,pos_list = self.translate_codon(seq[pos-1:],pos)
        end_pos = int((pos+2)/3) + len(aa_list)
        aa_pos_list = [i for i in range(int((pos+2)/3),end_pos)]
        cds_codon_list,cds_aa_list = self.extract_cds_seq(pos_list)
        cdna_codon_lines = self._split_func(cdna_list)
        cdna_aa_lines = self._split_func(aa_list)
        cds_codon_lines = self._split_func(cds_codon_list)
        cds_aa_lines = self._split_func(cds_aa_list)
        pos_lines = self._split_func(pos_list)
        aa_pos_lines = self._split_func(aa_pos_list)
        for pos,aa_pos,cds_seq,cds_aa,cdna_seq,cdna_aa in zip(pos_lines,aa_pos_lines,cds_codon_lines,cds_aa_lines,cdna_codon_lines,cdna_aa_lines):
            cds_seq_line = "".join([" {codon} ".format(codon=s) for s in cds_seq])

            cds_aa_line = "".join(["{a:^5}".format(a=a) for a in cds_aa])
            cdna_seq_line = "".join([" {codon} ".format(codon=s) for s in cdna_seq])
            cdna_aa_line = "".join(["{a:^5}".format(a=a) for a in cdna_aa])
            order = "".join(["{:^5}".format(p) for p in pos])
            aa_order = "".join(["{:^5}".format(p) for p in aa_pos])
            print(order)
            print(aa_order)
            print(cds_seq_line)
            print(cds_aa_line)
            print(cdna_seq_line)
            print(cdna_aa_line)
            print('-'*len(cds_seq_line))

    def translate_codon(self,seq,pos):
        """
        translate the codon to amino acid until terminal codon
        """
        idx = 0
        cdna_list = []
        aa_list = []
        pos_list = []
        while True:
            start = idx*3
            end = (idx+1)*3
            pos_list.append(pos+start)
            # print(start,end)
            try:
                codon = seq[start:end]
                # print(repr(codon))
                if len(codon) != 3:
                    cdna_list.append(seq[start:])
                    aa_list.append('')
                    break
            except IndexError:
                cdna_list.append(seq[start:])
                aa_list.append('')
                break
            aa = codon_tab[codon.upper()][0]
            cdna_list.append(codon)
            aa_list.append(aa)
            if aa == '*':
                break
            idx +=1
        # print(pos_list)
        return cdna_list,aa_list,pos_list

    def split_aa(self,aa):
        """
        split amino acid by num of each line
        """
        aa_lines = self._split_func(aa)
        if len(aa_lines[-1]) ==self.num_of_each_line:
            aa_lines.append(["*"])
        else:
            aa_lines[-1] += "*"
        return aa_lines

    def _split_func(self,coding_list):
        """
        split coding list or aa string by num of each line
        """
        seq_lines = []
        num_of_lines = int(len(coding_list)/self.num_of_each_line)
        if num_of_lines == 0:
            seq_lines = [coding_list]
        else:
            start = 0
            for n in range(num_of_lines):
                end = (n+1)*self.num_of_each_line
                seq_lines.append(coding_list[start:end])
                start = end
            
            left = coding_list[num_of_lines*self.num_of_each_line:]
            if left:
                seq_lines.append(left)

        return seq_lines

    def split_cds_seq(self,cds_seq):
        """
        split cds seq by 3 codon
        """
        coding_list = []
        print(len(cds_seq))
        cds_fold = int(len(cds_seq)/3)
        print(cds_fold)
        for i in range(cds_fold):
            start = i*3
            end = (i+1)*3
            coding_list.append(cds_seq[start:end])

        seq_lines = self._split_func(coding_list)
        return seq_lines
        
    def show_variants(self,seq_lines,aa_lines):
        """
        Show amino acid change by pretty print pattern
        """
        for idx,(seq,aa) in enumerate(zip(seq_lines,aa_lines)):
            seq_line = "".join([" {codon} ".format(codon=s) for s in seq])
            aa_line = "".join(["{a:^5}".format(a=a) for a in aa])
            order = "".join(["{:^5}".format(i+1) for i in range(self.num_of_each_line*idx,self.num_of_each_line*(idx+1))])
            print(order)
            print(seq_line)
            print(aa_line)
            print('-'*len(seq_line))
    
    def main(self):
        seq_lines = self.split_cds_seq(self.cds_seq)
        aa_lines = self.split_aa(self.translation)
        self.show_variants(seq_lines,aa_lines)

def test_main(NM,cdna):
    """
    test for cdna
    """
    xml_tree = get_NM_info(NM)
    cds_seq,translation,utr3 = parse_xml(xml_tree)
    shower = ShowVar(cdna,cds_seq,translation,utr3)
    shower.show_cdna_main()

if __name__ == "__main__":
    # xml_tree = get_NM_info('NM_001127208')
    # cds_seq,translation,utr3 = parse_xml(xml_tree)
    # shower = ShowVar('c.5436delC',cds_seq,translation,utr3)
    # shower.show_cdna_main()
    # shower = ShowVar('c.3309_3310del',cds_seq,translation,utr3)
    # shower.show_cdna_main()
    # shower = ShowVar('c.5668_5669del',cds_seq,translation,utr3)
    # shower.show_cdna_main()
    # shower.main()
    # test_main('NM_001127208','c.5436delC')
    # test_main('NM_001127208','c.3309_3310del')
    test_main('NM_001127208','c.5668_5669del')
    # test_main('NM_015338','c.1924_1928del')
    # test_main('NM_015338','c.1927dupG')